# create by fanfan on 2020/3/31 0031
import numpy as np
import os
import tensorflow as tf
import keras
from keras import backend as K
from keras import layers
from keras.callbacks import Callback
from tensorflow.contrib.learn.python.learn.datasets import mnist

if K.backend() != 'tensorflow':
    raise RuntimeError('This example can only run with the '
                       'TensorFlow backend, '
                       'because it requires TFRecords, which '
                       'are not supported on other platforms.')

class EvaluateInputTensor(Callback):
    def __init__(self,model,steps,metrics_prefix='val',verbose=1):
        # parameter of callbacks passed during initialization
        # pass evalation mode directly
        super(EvaluateInputTensor,self).__init__()
        self.val_model = model
        self.num_steps = steps
        self.verbose = verbose
        self.metrics_prefix = metrics_prefix

    def on_epoch_end(self,epoh,logs={}):
        self.val_model.set_weights(self.model.get_weights())
        results = self.val_model.evaluate(None,None,steps=int(self.num_steps),
                                          verbose = self.verbose)
        metrics_str = "\n"
        for result,name in zip(results,self.val_model.metrics_names):
            metric_name = self.metrics_prefix + "_" + name
            logs[metric_name] = result
            if self.verbose > 0:
                metrics_str = metrics_str + metric_name+":"+str(result)+ " "
        if self.verbose > 0:
            print(metrics_str)

def cnn_layers(x_train_input):
    x = layers.Conv2D(32,(3,3),
                      activation='relu',padding='valid')(x_train_input)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(64,(3,3),activation='relu')(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Flatten()(x)
    x = layers.Dense(512,activation='relu')(x)
    x = layers.Dropout(0.5)(x)
    x_train_out = layers.Dense(num_classes,
                               activation='softmax',
                               name='x_train_out')(x)
    return x_train_out

sess = K.get_session()
batch_size = 100
batch_shape = (batch_size,28,28,1)
epochs = 5
num_classes = 10

capacity = 10000

min_after_dequeue = 3000

enqueue_many = True
cache_dir = os.path.expanduser(
    os.path.join('~', '.keras', 'datasets', 'MNIST-data'))
data = mnist.read_data_sets(cache_dir, validation_size=0)

x_train_batch,y_train_batch = tf.train.shuffle_batch(
    tensors=[data.train.images,data.train.labels.astype(np.int32)],
    batch_size = batch_size,
    capacity=capacity,
    min_after_dequeue=min_after_dequeue,
    enqueue_many=enqueue_many,
    num_threads=8
)

x_train_batch = tf.cast(x_train_batch,tf.float32)
x_train_batch = tf.reshape(x_train_batch,shape=batch_shape)

y_train_batch = tf.cast(y_train_batch,tf.int32)
y_train_batch = tf.one_hot(y_train_batch,num_classes)

x_batch_shape = x_train_batch.get_shape().as_list()
y_batch_shape = y_train_batch.get_shape().as_list()

model_input = layers.Input(tensor=x_train_batch)
model_output = cnn_layers(model_input)
train_model = keras.models.Model(inputs=model_input,outputs=model_output)

# Pass the target tensor `y_train_batch` to `compile`
# via the `target_tensors` keyword argument:
train_model.compile(optimizer=keras.optimizers.RMSprop(lr=2e-3, decay=1e-5),
                    loss='categorical_crossentropy',
                    metrics=['accuracy'],
                    target_tensors=[y_train_batch])
train_model.summary()

x_test_batch, y_test_batch = tf.train.batch(
    tensors=[data.test.images, data.test.labels.astype(np.int32)],
    batch_size=batch_size,
    capacity=capacity,
    enqueue_many=enqueue_many,
    num_threads=8)

# Create a separate test model
# to perform validation during training
x_test_batch = tf.cast(x_test_batch, tf.float32)
x_test_batch = tf.reshape(x_test_batch, shape=batch_shape)

y_test_batch = tf.cast(y_test_batch, tf.int32)
y_test_batch = tf.one_hot(y_test_batch, num_classes)

x_test_batch_shape = x_test_batch.get_shape().as_list()
y_test_batch_shape = y_test_batch.get_shape().as_list()

test_model_input = layers.Input(tensor=x_test_batch)
test_model_output = cnn_layers(test_model_input)
test_model = keras.models.Model(inputs=test_model_input, outputs=test_model_output)

# Pass the target tensor `y_test_batch` to `compile`
# via the `target_tensors` keyword argument:
test_model.compile(optimizer=keras.optimizers.RMSprop(lr=2e-3, decay=1e-5),
                   loss='categorical_crossentropy',
                   metrics=['accuracy'],
                   target_tensors=[y_test_batch])

coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess,coord)
train_model.fit(
    epochs=epochs,
    steps_per_epoch=int(np.ceil(data.train.num_examples/float(batch_size))),
    callbacks=[EvaluateInputTensor(test_model,steps=100)]
)
# Save the model weights.
train_model.save_weights('saved_wt.h5')

# Clean up the TF session.
coord.request_stop()
coord.join(threads)
K.clear_session()

# Second Session to test loading trained model without tensors
x_test = np.reshape(data.test.images, (data.test.images.shape[0], 28, 28, 1))
y_test = data.test.labels
x_test_inp = layers.Input(shape=(x_test.shape[1:]))
test_out = cnn_layers(x_test_inp)
test_model = keras.models.Model(inputs=x_test_inp, outputs=test_out)

test_model.load_weights('saved_wt.h5')
test_model.compile(optimizer='rmsprop',
                   loss='categorical_crossentropy',
                   metrics=['accuracy'])
test_model.summary()

loss, acc = test_model.evaluate(x_test,
                                keras.utils.to_categorical(y_test),
                                batch_size=batch_size)
print('\nTest accuracy: {0}'.format(acc))



